from datetime import datetime
from collections import defaultdict
from sqlalchemy import func, select, exc, and_
from app.config.db import DBUtils
from app.model.FKListModel import fklist
from app.model.SKListModel import sklist
from app.model.BusinessInfoModel import BusinessInfo

# 图谱数据结构处理
def process_graph(result_list, comb, lx):
    group_fkje_nd = defaultdict(int)
    label_set = set()
    total_amount = 0
    gs_map = {
        '深研院':'深圳研究院',
        '集团公司':'集团总部',
        '航天科技财务有限责任公司':'财务公司',
        '档案馆':'航天档案馆',
        '系统院':'系统工程院',
        '资源卫星':'资源卫星中心'
    }
    gs_id = {11989570, 11989569, 11989571}
    for row in result_list:
        ys_gs = row['gs']
        if row['ID'] in gs_id:
            gs = '七院'
        else:
            gs = gs_map.get(ys_gs, ys_gs)
        key_nd = (row['cydw'], gs, row['ID'])
        group_fkje_nd[key_nd] += row['total']
        if gs not in label_set:
            label_set.add(gs)
        total_amount += row['total']
    node = [{
        'id': lx + str(ID),
        'label': cydw,
        'nodeType': 'cydw',
        'amount': str(round(total / 10000, 2)) if total >= 50 or total == 0 else '0.01',
        'combo': lx + gs
    } for (cydw, gs, ID), total in group_fkje_nd.items()]

    combos = [{
        'id': lx + gs,
        'label': gs,
        'combo': comb
    } for gs in label_set]
    return node, combos, total_amount

class BelongingFlowService:
    def __init__(self):
        self.db_utils = DBUtils()
        self.current_year = datetime.now().year

    def select_belonging_flow_info(self, session, company, is_cydw, year):
        if int(is_cydw) in range(1,4):
            return None
        combos = [
            {'id': 'cb5', 'label': '上游归属'},
            {'id': 'cb6', 'label': '下游归属'}
        ]
        # {'id': 'node4', 'label': '上游供应商', 'nodeType': 'wldw', 'combo': 'cb5'},
        # {'id': 'node6', 'label': '下游客户', 'nodeType': 'wldw', 'combo': 'cb6'}
        company_node = [{'id': 'node5', 'label': company, 'nodeType': 'wldw'}]
        sk_edge = [{'source': 'cb6', 'target': 'node5', 'elx': 'SK'}]
        fk_edge = [{'source': 'node5', 'target': 'cb5', 'elx': 'FK'}]

        # if start_date and end_date:
        #     condition_fk = fklist.SCRQ.between(start_date, end_date)
        #     condition_sk = sklist.SCRQ.between(start_date, end_date)
        # else:
        condition_fk = fklist.JYND == year
        condition_sk = sklist.JYND == year
        try:
            # upstream_query = (
            #     select(sklist.SKGS,
            #            func.sum(sklist.FKJE).label('total'))
            #     .where(and_(sklist.FKZHMC == company,
            #                 condition_sk, sk_dif_name))
            #     .group_by(sklist.SKGS)
            # )
            # upstream_list = session.execute(upstream_query).fetchall()
            upstream_query = (
                select(sklist.SKZHMC.label('cydw'),
                       sklist.SKGS.label('gs'),
                       BusinessInfo.ID,
                       func.sum(sklist.FKJE).label('total'))
                .join(BusinessInfo, BusinessInfo.GSMC == sklist.SKZHMC)
                .where(and_(sklist.FKZHMC == company, condition_sk, sklist.SKGS.isnot(None)))
                .group_by(sklist.SKZHMC, sklist.SKGS, BusinessInfo.ID)
                .order_by(func.sum(sklist.FKJE).desc())
                .limit(5)
            )
            # print(upstream_query.compile(compile_kwargs={"literal_binds": True}))
            upstream_list = session.execute(upstream_query).mappings()

            downstream_query = (
                select(fklist.FKZHMC.label('cydw'),
                       fklist.FKGS.label('gs'),
                       BusinessInfo.ID,
                       func.sum(fklist.FKJE).label('total'))
                .join(BusinessInfo, BusinessInfo.GSMC == fklist.FKZHMC)
                .where(and_(fklist.SKZHMC == company, condition_fk, fklist.FKGS.isnot(None)))
                .group_by(fklist.FKZHMC, fklist.FKGS, BusinessInfo.ID)
                .order_by(func.sum(fklist.FKJE).desc())
                .limit(5)
            )
            # print(downstream_query.compile(compile_kwargs={"literal_binds": True}))
            downstream_list = session.execute(downstream_query).mappings()
            # downstream_query = (
            #     select(fklist.FKGS,
            #            func.sum(fklist.FKJE).label('total'))
            #     .where(and_(fklist.SKZHMC == company,
            #                 condition_fk, fk_dif_name))
            #     .group_by(fklist.FKGS)
            # )
            # downstream_list = session.execute(downstream_query).fetchall()

            if upstream_list:
                upstream_nodes, upstream_combos, up_total = process_graph(upstream_list, 'cb5', 'up')
                fk_edge[0]['label'] = str(round(up_total / 10000, 2)) if up_total >= 50 or up_total == 0 else '0.01'
            else:
                upstream_nodes = []
                upstream_combos = []
                fk_edge = [{'source': 'node5', 'target': 'cb5', 'elx': 'FK', 'label': 0}]

            if downstream_list:
                downstream_nodes, downstream_combos, down_total = process_graph(downstream_list, 'cb6', 'down')
                sk_edge[0]['label'] = str(round(down_total / 10000, 2)) if down_total >= 50 or down_total == 0 else '0.01'
            else:
                downstream_nodes = []
                downstream_combos = []
                sk_edge = [{'source': 'cb6', 'target': 'node5', 'elx': 'SK', 'label': 0}]

            nodes = upstream_nodes + downstream_nodes + company_node
            edges = sk_edge + fk_edge
            combs = combos + upstream_combos + downstream_combos
            result = [{
                'nodes': nodes,
                'edges': edges,
                'combos': combs,
            }]
        except exc.SQLAlchemyError as e:
            error_response = [{'error': str(e)}]
            try:
                self.db_utils.rollback()
            except Exception as rollback_error:
                print(f"Failed to rollback transaction: {rollback_error}")
            return error_response
        finally:
            try:
                self.db_utils.remove_session()
            except Exception as remove_error:
                print(f"Failed to rollback transaction: {remove_error}")

        return result

        #     if upstream_list:
        #         up_id = 1
        #         upstream_nodes = []
        #         upstream_edges = []
        #         for row in upstream_list:
        #             strid = 'nd' + str(up_id)
        #             if row.total >= 50:
        #                 str_total = str(round(row.total / 10000, 2))
        #             else:
        #                 str_total = '0.01'
        #             upstream_node = {
        #                 'id': strid,
        #                 'label': row.SKGS,
        #                 'nodeType': 'gs',
        #                 'amount': str_total,
        #                 'combo': 'cb5'}
        #             upstream_nodes.append(upstream_node)
        #             upstream_edge = {
        #                 'source': strid,
        #                 'target': 'node4',
        #                 'label': str_total,
        #                 'elx': 'gs'}
        #             upstream_edges.append(upstream_edge)
        #             up_id += 1
        #
        #         up_total = sum(float(row['label']) for row in upstream_edges)
        #         fk_edge[0]['label'] = str(round(up_total, 2))
        #     else:
        #         upstream_nodes = []
        #         upstream_edges = []
        #         fk_edge = []
        #
        #     if downstream_list:
        #         down_id = 1000
        #         downstream_nodes = []
        #         downstream_edges = []
        #         for row in downstream_list:
        #             strid = 'nd' + str(down_id)
        #             if row.total >= 50:
        #                 str_total = str(round(row.total / 10000, 2))
        #             else:
        #                 str_total = '0.01'
        #             downstream_node = {
        #                 'id': strid,
        #                 'label': row.FKGS,
        #                 'nodeType': 'gs',
        #                 'amount': str_total,
        #                 'combo': 'cb6'}
        #             downstream_nodes.append(downstream_node)
        #             downstream_edge = {
        #                 'source': strid,
        #                 'target': 'node6',
        #                 'label': str_total,
        #                 'elx': 'gs'}
        #             downstream_edges.append(downstream_edge)
        #             down_id += 1
        #
        #         down_total = sum(float(row['label']) for row in downstream_edges)
        #         sk_edge[0]['label'] = str(round(down_total, 2))
        #     else:
        #         downstream_nodes = []
        #         downstream_edges = []
        #         sk_edge = []
        #
        #     nodes = upstream_nodes + downstream_nodes + company_node
        #     edges = upstream_edges + downstream_edges + sk_edge + fk_edge
        #     result = [{
        #         'nodes': nodes,
        #         'edges': edges,
        #         'combos': combos,
        #     }]
        # except exc.SQLAlchemyError as e:
        #     error_response = [{'error': str(e)}]
        #     try:
        #         self.db_utils.rollback()
        #     except Exception as rollback_error:
        #         print(f"Failed to rollback transaction: {rollback_error}")
        #     return error_response
        # finally:
        #     try:
        #         self.db_utils.remove_session()
        #     except Exception as remove_error:
        #         print(f"Failed to rollback transaction: {remove_error}")
        #
        # return result
